# -*- coding: utf-8 -*-
# 根据传入的图片路径，生成图片向量
# 模型地址：https://www.modelscope.cn/models/iic/multi-modal_clip-vit-base-patch16_zh/summary/
# @Time    : 2024/12/8 下午10:50
# @Author  : 赵鑫源
# @FileName: generate_vectors.py
# @Software: PyCharm
import io

import numpy as np
import requests
from PIL import Image
import global_vars  # 全局变量
from mysql_helper import MySQLHelper
import my_logging
import ulid
import json

mysql_helper = MySQLHelper()  # 初始化
logging = my_logging.LogUtils().get_log()


def update_index_to_ulid(img_id):
    """
        更新 index_to_ulid 映射
        img_id: 图片 id
    """
    existing_ids = set(global_vars.index_to_ulid.values())
    if img_id not in existing_ids:
        idx = len(global_vars.index_to_ulid)  # 使用当前长度作为索引
        global_vars.index_to_ulid[idx] = img_id


def add_update_vectors(img_object):
    """
        新增或者修改图片向量
    """
    global_vars.global_i2key[img_object['id']] = img_object

    # 更新数据库中的向量
    result = mysql_helper.execute_query('select img_id from t_images_vectors where img_id = %s', (str(img_object['id']),))
    if result:  # 如果存在，则更新
        mysql_helper.execute_update('update t_images_vectors set img_vectors = %s where img_id = %s',
                                    (str(img_object['imgVectors']), str(img_object['id'])))
        # 更新 index_to_ulid
        update_index_to_ulid(img_object['id'])
        # 更新索引
        rebuild_vectors()
    else:  # 如果不存在，则新增
        mysql_helper.execute_insert('insert into t_images_vectors (id,img_id,img_vectors) values (%s, %s, %s)',
                                    (str(ulid.new()), str(img_object['id']), str(img_object['imgVectors'])))

        # 更新 index_to_ulid
        update_index_to_ulid(img_object['id'])
        # 添加新向量到 Faiss 索引
        _vec = np.array([float(x) for x in img_object['imgVectors'].split(',')]).reshape(1, 512)
        global_vars.global_index.add(_vec)


def update_img_vectors_status(id, status):
    """
        更新图片向量状态
        id: 图片id
    """
    mysql_helper.execute_update('update t_images_info set img_vectors_status = %s where id = %s',
                                (status, str(id)))


def del_vectors(img_id):
    """
        删除内存中的图片向量
        img_id: 图片id
        is_memory: 是否更新内存中的向量，默认为True
    """
    try:
        # 删除 global_i2key 中的记录
        if img_id in global_vars.global_i2key:
            del global_vars.global_i2key[img_id]

        # 删除 index_to_ulid 中的记录
        for k, v in list(global_vars.index_to_ulid.items()):  # 使用 list 避免迭代时修改字典
            if v == img_id:
                del global_vars.index_to_ulid[k]
                break
        # 更新索引
        rebuild_vectors()
    except Exception as e:
        logging.error('删除内存中的向量失败', e)


def rebuild_vectors():
    """
        重建索引
    """
    # 重建索引
    vectors = []
    for v in global_vars.global_i2key.values():
        _vec = np.array([float(x) for x in v['imgVectors'].split(',')]).reshape(1, 512)
        vectors.append(_vec)
    if vectors:
        global_vars.global_index.reset()
        global_vars.global_index.add(np.vstack(vectors))


def format_record(img_embedding):
    """
       格式化数据
       img_embedding: 图片向量
    """
    return ",".join([str(num) for num in img_embedding])


def build_vectors(img_object):
    """
        构建索引
        img_object: 图片信息对象
    """
    try:
        # 如果图片被删除，则删除内存中的向量并返回
        if img_object['flag'] == 'DELETED':
            del_vectors(img_object['id'])
            return json.dumps({"code": 200, "msg": "TBIR删除图片成功"})

        # 根据图片路径获取图片特征
        response = requests.get(img_object['bucketUrl'] + img_object['imgPngPath'])
        response.raise_for_status()  # 检查请求是否成功
        img_data = response.content
        # 加载图片
        img_path = Image.open(io.BytesIO(img_data))
        # 2D Tensor, [图片数, 特征维度]
        img_embedding = global_vars.global_model.forward({'img': img_path})['img_embedding']
        # 格式化数据
        current_vector = format_record(img_embedding[0].numpy())
        # 更新数据库中的vectors字段
        img_object['imgVectors'] = current_vector

        # 状态是新增或者更新
        if img_object['flag'] == 'ADD' or img_object['flag'] == 'UPDATE':
            add_update_vectors(img_object)
            update_img_vectors_status(img_object['id'], 'SUCCESS')
        logging.info('解析图片向量成功')
        return json.dumps({"code": 200, "msg": "TBIR解析图片成功"})
    except Exception as e:
        logging.error('图片解析异常', e)
        update_img_vectors_status(img_object['id'], 'FAIL')
        return json.dumps({"code": 200, "msg": "TBIR解析图片失败"})
